from math import sin, cos, pi
import glm
import OpenGL.GL as gl
import numpy as np
import pygl

class Sphere:
    def __init__(self, radius=1.0, u_segments=32, v_segments=32) -> None:
        number_vertex = (u_segments + 1) * (v_segments + 1)
        self.number_indices = u_segments * v_segments * 6
        self.vertices = np.empty((number_vertex, 8), dtype=gl.GLfloat)
        self.indices = np.empty((self.number_indices), dtype=gl.GLuint)
        for v in range(v_segments + 1):
            for u in range(u_segments + 1):
                temp_u = u / float(u_segments)
                temp_v = v / float(v_segments)
                theta = temp_u * pi
                phi = temp_v * pi * 2
                index = u + (u_segments + 1) * v
                x = sin(theta) * cos(phi)
                y = sin(theta) * sin(phi)
                z = cos(theta)
                # 位置
                self.vertices[index, 0] = x * radius
                self.vertices[index, 1] = y * radius
                self.vertices[index, 2] = z * radius
                # 法线向量
                self.vertices[index, 3] = x
                self.vertices[index, 4] = y
                self.vertices[index, 5] = z
                # 纹理坐标
                self.vertices[index, 6] = temp_v
                self.vertices[index, 7] = temp_u
        index = 0
        for v in range(v_segments):
            for u in range(u_segments):
                vindex = u + (u_segments + 1) * v
                self.indices[index + 0] = vindex
                self.indices[index + 1] = vindex + 1
                self.indices[index + 2] = vindex + 1 + (u_segments + 1)
                self.indices[index + 3] = vindex
                self.indices[index + 4] = vindex + 1 + (u_segments + 1)
                self.indices[index + 5] = vindex + (u_segments + 1)
                index += 6
        self.vao = pygl.VertexArrayObject()
        self.vbo = pygl.VertexBufferObject(self.vertices)
        self.ebo = pygl.ElementBufferObject(self.indices)
        attribute_position = pygl.VertexAttribute("position", 0, 3, 
                                            gl.GL_FLOAT, gl.GL_FALSE, 0)
        attribute_normal = pygl.VertexAttribute("normal", 1, 3, gl.GL_FLOAT, 
                                    gl.GL_FALSE, 3 * gl.sizeof(gl.GLfloat))
        attribute_tex_coord = pygl.VertexAttribute("tex_coord", 2, 2, 
                gl.GL_FLOAT, gl.GL_FALSE, 6 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexBuffer(self.vbo, 0, 0, 8 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexAttribute(0, attribute_position)
        self.vao.setVertexAttribute(0, attribute_normal)
        self.vao.setVertexAttribute(0, attribute_tex_coord)
        self.vao.setElementBuffer(self.ebo)
        self.rotation = glm.vec3(0.0, 0.0, 0.0)
        self.position = glm.vec3(0.0, 0.0, 0.0)
        self.scale = glm.vec3(1.0, 1.0, 1.0)
        self.updateModelMatrix()

    def render(self)->None:
        self.vao.bind()
        gl.glDrawElements(gl.GL_TRIANGLES, self.number_indices, 
                          gl.GL_UNSIGNED_INT, None)

    def delete(self)->None:
        self.vao.delete()
        self.vbo.delete()
        self.ebo.delete()

    def updateModelMatrix(self)->None:
        scale_matrix = glm.scale(glm.mat4(1.0), self.scale)
        rotation_matrix = glm.rotate(glm.mat4(1.0), glm.radians(self.rotation.x), 
                                     glm.vec3(1.0, 0.0, 0.0))
        rotation_matrix = glm.rotate(rotation_matrix, 
                        glm.radians(self.rotation.y), glm.vec3(0.0, 1.0, 0.0))
        rotation_matrix = glm.rotate(rotation_matrix, 
                        glm.radians(self.rotation.z), glm.vec3(0.0, 0.0, 1.0))
        translation_matrix = glm.translate(glm.mat4(1.0), self.position)
        self.model_matrix = translation_matrix * rotation_matrix * scale_matrix

class InstancedSphere(Sphere):
    def __init__(self, instances:np.ndarray, radius=1.0, u_segments=32, 
                 v_segments=32) -> None:
        super().__init__(radius, u_segments, v_segments)
        self.instanced_vbo = pygl.VertexBufferObject(instances)
        self.number_instance = instances.shape[0]
        instance_position = pygl.VertexAttribute("instance_position", 3, 3, 
                                                 gl.GL_FLOAT, gl.GL_FALSE, 0)
        instance_color = pygl.VertexAttribute("instance_color", 4, 3, 
                        gl.GL_FLOAT, gl.GL_FALSE, 3 * gl.sizeof(gl.GLfloat))
        instance_radius = pygl.VertexAttribute("instance_radius", 5, 1, 
                        gl.GL_FLOAT, gl.GL_FALSE, 6 * gl.sizeof(gl.GLfloat))
        self.vao.setVertexBuffer(self.instanced_vbo, 1, 0, 7 * 
                                 gl.sizeof(gl.GLfloat))
        self.vao.setVertexAttribute(1, instance_position)
        self.vao.setVertexAttribute(1, instance_color)
        self.vao.setVertexAttribute(1, instance_radius)
        self.vao.setBindingDivisor(1, 1)

    def render(self)->None:
        self.vao.bind()
        gl.glDrawElementsInstanced(gl.GL_TRIANGLES, self.number_indices, 
                                   gl.GL_UNSIGNED_INT, None,  self.number_instance)
